import sys
import os
import json
import requests
import time
from xhs import XhsClient, DataFetchError, help
from utils import sign, read_cookie_from_file


# def extract_ids_from_file(file_path):
#     # 打开并读取JSON文件
#     with open(file_path, 'r', encoding='utf-8') as file:
#         data = json.load(file)

#     # 提取items中的id字段
#     ids = [item['id'] for item in data.get('items', [])]

#     return ids

def extract_ids_and_tokens_from_file(file_path):
    # 打开并读取JSON文件
    with open(file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)

    # 提取每个元素的id和xsec_token字段，并配对
    pairs = []
    for item in data:
        if 'id' in item and 'xsec_token' in item:
            pairs.append({'id': item['id'], 'xsec_token': item['xsec_token']})

    # 返回包含id和token配对的列表
    return pairs



def download_images(img_urls, folder_path):
    for idx, url in enumerate(img_urls):
        try:
            response = requests.get(url)
            response.raise_for_status()  # 如果请求失败，抛出异常
            image_path = os.path.join(folder_path, f"{idx + 1}.png")
            with open(image_path, 'wb') as img_file:
                img_file.write(response.content)
            print(f"Saved image {idx + 1} to {image_path}")
        except requests.exceptions.RequestException as e:
            print(f"Failed to download image {url}: {e}")


def save_note_to_folder(note, img_urls, base_folder):
    # 创建一个以note_id命名的文件夹
    folder_path = os.path.join(base_folder, note['note_id'])
    os.makedirs(folder_path, exist_ok=True)

    # 创建一个文本文件，保存笔记内容
    file_path = os.path.join(folder_path, 'note_content.txt')
    with open(file_path, 'w', encoding='utf-8') as file:
        file.write(f"Title: {note['title']}\n")
        file.write(f"Note ID: {note['note_id']}\n")
        file.write(f"Description: {note['desc']}\n")

    # 写入note_info到文件
    note_info_path = os.path.join(folder_path, 'note_info.json')
    with open(note_info_path, 'w', encoding='utf-8') as f:
        json.dump(note, f, ensure_ascii=False, indent=4)

    # 写入图片URL到文件
    img_urls_path = os.path.join(folder_path, 'img_urls.json')
    with open(img_urls_path, 'w', encoding='utf-8') as f:
        json.dump(img_urls, f, ensure_ascii=False, indent=4)

    # 下载并保存图片
    download_images(img_urls, folder_path)

    print(f"Note content and image URLs saved to: {folder_path}")

def run_file2(keyword):
    # 关键字路径
    base_folder = os.path.join(os.getcwd(), keyword)
    print(base_folder, '==== base folder')
    file_path = os.path.join(base_folder, 'data.json')


    # 读取cookie
    cookie = read_cookie_from_file()

    xhs_client = XhsClient(cookie, sign=sign)

    # 提取ID列表
    items = extract_ids_and_tokens_from_file(file_path)
    for item in items:
        for _ in range(10):
            try:
                time.sleep(1)
                note_info = xhs_client.get_note_by_id(item['id'], 
                                                      item['xsec_token'])
                img_urls = help.get_imgs_url_from_note(note_info)

                save_note_to_folder(note_info, img_urls, base_folder)
                break  # 如果获取成功，则跳出重试循环
            except DataFetchError as e:
                print(e)
                print("失败重试一下下")

if __name__ == '__main__':
    if len(sys.argv) > 1:
        keyword = sys.argv[1]  # 从命令行获取关键字
    else:
        keyword = input("请输入关键字: ")  # 如果没有传递关键字，提示用户输入

    run_file2(keyword)
